import pickle
import socket
from . import io
from . import config
from .connection import Connection
from . import logger
import threading

__all__ = ["ImServer", "ImManager"]

logger = logger.get_logger(__package__)

class Command:
    _id = 1

    @classmethod
    def id(cls):
        val = cls._id
        cls._id += 1
        return val

    def __init__(self, command, args):
        self.command = command
        self.args = args
        self.rid = self.id()
        self._result = None
        self._lock = threading.Lock()
        self._lock.acquire()

    def __str__(self):
        return self.command

    @property
    def result(self):
        return self._result

    def wait_done(self, timeout=-1):
        return self._lock.acquire(timeout=timeout)

    def done(self, result):
        self._result = result
        self._lock.release()


class ImManager(io.Client):
    def __init__(self, address):
        super(ImManager, self).__init__(address)
        self._id_command_map = {}

    _cmd_handler_map = dict()

    @classmethod
    def register_command(cls, command, handler):
        cls._cmd_handler_map[command] = handler

    @classmethod
    def unregister_command(cls, command):
        del cls._cmd_handler_map[command]

    def process_command(self, command, args):
        return self._cmd_handler_map[command](self, args)

    def on_connect(self, sock, remote):
        super(ImManager, self).on_connect(sock, remote)

    def on_disconnect(self, sock):
        super(ImManager, self).on_disconnect(sock)

    def recv_data(self, sock):
        return ImServer.recv(sock)

    def send_data(self, sock, data):
        return ImServer.send(sock, data)

    def on_received(self, sock, data):
        try:
            rid, errcode, args = pickle.loads(data)
            if rid in self._id_command_map:
                cmd = self._id_command_map[rid]
                cmd.done((errcode, data))
        except Exception as e:
            logger.error(str(e))

    def on_sent(self, sock, data, sz_data):
        pass

    def _send_command(self, cmd, args=None, timeout=-1):
        command = Command(cmd, args)
        data = pickle.dumps((command.rid, command.command, command.args))
        self._id_command_map[command.rid] = command
        try:
            logger.info("The command {} has been submitted".format(cmd))
            self.send_in_queue(self.sock, data)
            if command.wait_done(timeout=timeout):
                return command.result
        finally:
            del self._id_command_map[command.rid]


class ImServer(io.Server):
    def __init__(self, manager=None):
        super(ImServer, self).__init__()
        self.manager = manager
        self.config = config

    @property
    def manager(self):
        return self._manager

    @manager.setter
    def manager(self, manager):
        self._manager = manager
        if manager:
            setattr(self._manager, "server", self)

    @staticmethod
    def send(sock, data):
        data = (len(data)).to_bytes(5, byteorder='big') + data
        data_len = sock.send(data)
        logger.info("Sent {} bytes: {}".format(len(data)-5, data[5:]))
        return data_len == len(data)

    @staticmethod
    def recv(sock):
        try:
            data = sock.recv(5)
            if not data:
                raise ConnectionResetError
            data_len = int.from_bytes(data, byteorder='big')
            data = sock.recv(data_len)
            logger.info("Received {} bytes: {}".format(data_len, data))
            if len(data) != data_len:
                return None
            return data
        except ConnectionResetError:
            return None

    def get_sockets(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((self.config.bind_address[0], self.config.bind_address[1]))
        sock.listen(self.config.backlog)
        yield sock

        if self.manager:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.bind((self.config.bind_address[0], self.config.bind_address[1] + 1))
            sock.listen(self.config.backlog)
            setattr(self.manager, "_sock", sock)
            yield sock

    def is_manager(self, sock):
        assert isinstance(sock, socket.socket)
        if self.manager and hasattr(self.manager, "_sock") and getattr(self.manager, "_sock", sock):
            return True

        return False

    def setsockopt(self, sock):
        assert isinstance(sock, socket.socket)
        if self.config.snd_buf_size:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self.config.snd_buf_size)
        if self.config.rcv_buf_size:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self.config.snd_buf_size)

    def on_connect(self, client, server):
        sock = super(ImServer, self).on_connect(client, server)
        Connection(self, client)
        return sock

    def on_disconnect(self, client):
        super(ImServer, self).on_disconnect(client)
        conn = Connection.get_connection(client.fileno())
        conn.close()

    def send_data(self, sock, data):
        return ImServer.send(sock, data)

    def recv_data(self, sock):
        return ImServer.recv(sock)

    def on_received(self, sock, data):
        super(ImServer, self).on_received(sock, data)
        if self.is_manager(sock):
            rid, command, args = pickle.loads(data)
            logger.info("Processing the command {}".format(command))
            errcode, data = self.manager.process_command(command, args)
            self.send_in_queue(sock, pickle.dumps((rid, errcode, data)))
        else:
            conn = Connection.get_connection(sock.fileno())
            conn.process_data(data)

    def on_sent(self, sock, data, sz_send):
        super(ImServer, self).on_sent(sock, data, sz_send)

    def on_start(self):
        pass

    def start(self, **kwargs):
        def daemon(pid_file, chdir=None):
            import os, sys
            pid = os.fork()
            if pid > 0:
                raise SystemExit(0)

            if chdir:
                os.chdir(chdir)
            os.setsid()
            pid = os.fork()
            if pid > 0:
                raise SystemExit(0)

            # signal.signal(signal.SIGTERM, sig_handler)
            with open(pid_file, "w") as fp:
                fp.write("%d" % os.getpid())

        def set_log(log_dir):
            if log_dir:
                import sys
                sys.stdout = open(log_dir, "w+")

        for c in kwargs:
            setattr(self.config, c, kwargs[c])

        if self.config.daemon:
            daemon(self.config.pid_file, self.config.chdir)

        set_log(self.config.log_file)

        super(ImServer, self).start()


if __name__ == "__main__":
    p = ImServer()
    p.start()
